import pandas as pd
import ast
import pickle
import os

from tqdm import tqdm

from utils import search_data
from _base_plot import less_three_line_plot
from _base_plot import less_three_bar_plot
from _base_plot import plot_compare_bar_with_last_peirod
from _base_plot import plot_monthly_sesaonal_data
from _base_plot import plot_simple_bar_by_col
from _base_plot import plot_stack_col

from _llm import chat_with_theme


def build_args(save_name, plot_excel):
    # 从Excel文件中读取数据
    # 假设Excel文件中有两个工作表：'plot_data'和'filter_cols'
    """
    从Excel文件中读取数据，构建绘图函数的关键字参数。

    参数:
    - save_name (str): 保存名称，用于定位要读取的数据。
    - plot_excel (dict): 包含多个工作表的Excel文件的字典。
    """
    info = plot_excel['plot_data']
    info.index = info['save_name']
    filter_info = plot_excel['filter_cols']
    use_cols = plot_excel['use_cols']
    
    se = info.loc[save_name]
    kwargs = se.dropna().to_dict()
    
    if 'colors' in kwargs and isinstance(kwargs['colors'], str):
        kwargs['colors'] = ast.literal_eval(kwargs['colors'])
    
    if save_name in filter_info.index:
        filter_names = filter_info.loc[save_name].dropna().to_list()
        use_names = use_cols.loc[save_name].dropna().to_list()
        kwargs['filter_cols'] = filter_names
        kwargs['use_cols'] = use_names
    
    if 'show_period' in kwargs:
        kwargs['show_period'] = int(kwargs['show_period'])
    
    if 'ppt' in kwargs:
       del kwargs['ppt']
    
    return kwargs


def build_data(save_name, plot_excel):
    """
    从Excel文件中读取数据，构建绘图函数的关键字参数。
    """
    """
    从Excel文件中读取数据，构建绘图函数的关键字参数。

    参数:
    - save_name (str): 保存名称，用于定位要读取的数据。
    - plot_excel (dict): 包含多个工作表的Excel文件的字典。
    """
    info = plot_excel['plot_data']
    info.index = info['save_name']
    filter_info = plot_excel['filter_cols']
    use_cols = plot_excel['use_cols']
    
    se = info.loc[save_name]
    kwargs = se.dropna().to_dict()
    
    theme = kwargs['theme']
    print(theme, search_data(theme))
    data = pd.read_excel('cache/' + search_data(theme), index_col=0)
    
    data = data.dropna()
    
    if 'colors' in kwargs and isinstance(kwargs['colors'], str):
        kwargs['colors'] = ast.literal_eval(kwargs['colors'])
    
    if save_name in filter_info.index:
        filter_names = filter_info.loc[save_name].dropna().to_list()
        use_names = use_cols.loc[save_name].dropna().to_list()
        kwargs['filter_cols'] = filter_names
        
        filter_cols = kwargs['filter_cols']
        data = data.loc[:,filter_cols]
        
        kwargs['use_cols'] = use_names
        df = data.loc[:, filter_names]
    
    elif 'name' in kwargs:
        name = kwargs['name']
        df = data.loc[:, [name]]
    
    if 'show_period' in kwargs:
        show_period = kwargs['show_period']
        df = df.iloc[-int(show_period):,:]
        
    return df


def gen_img_data_chat(chat=True):

    """
    生成图像、数据和聊天文本的集合。
    """
    plot_selector = {
        'lineplot': less_three_line_plot,
        'barplot': less_three_bar_plot,
        'bar_compare_with_last_period': plot_compare_bar_with_last_peirod,
        'seasonal_lineplot': plot_monthly_sesaonal_data,
        'simple_bar_plot': plot_simple_bar_by_col,
        'bar_stack': plot_stack_col,
    }

    plot_excel = pd.read_excel(
        'meta\plot_args_use.xlsx', 
        sheet_name=None,
        index_col=0,
        )

    info = plot_excel['plot_data']
    info.index = info['save_name']

    img_collection = {} # 用于存储图像
    data_collection = {} # 用于存储dat
    chat_collection = {} # 用于存储chat

    for i in tqdm(range(len(info))):
        save_name = info.index[i]
        se = info.iloc[i,:]
    
        theme = se['theme']
        theme_collector = img_collection.setdefault(theme, {})
        theme_data_collector = data_collection.setdefault(theme, {}) # 用于存储data
        theme_chat_collector = chat_collection.setdefault(theme, {}) # 用于存储chat

        plot_func = plot_selector[se['plot_type']]
        kwargs = build_args(save_name, plot_excel)
        data_df = build_data(save_name, plot_excel) # 用于存储data
        
        print(kwargs)
        del kwargs['plot_type'] # 删除不需要的关键字参数
        
        print(plot_func.__name__, 'Plot>>>', kwargs)
        
        theme_collector[save_name] = plot_func(**kwargs)
        theme_data_collector[save_name] = data_df # 用于存储data
        
        if chat:
            chat_text = chat_with_theme(data_df, theme) # 用于存储chat
            theme_chat_collector[save_name] = chat_text # 用于存储chat
        else:
            print('no chat')

        
    collector = {
        'img_collection': img_collection,
        'data_collection': data_collection, # 用于存储data
        'chat_collection': chat_collection, # 用于存储chat
    }
    return collector


def load_plot_and_chat_data(load_cache=True, save_cache=True, chat=True):
    if not os.path.exists('cache\collector.pkl') or not load_cache:
        collector = gen_img_data_chat(chat=chat)
        print('loading cache')
        if save_cache:
            with open('cache\collector.pkl', 'wb') as f:
                pickle.dump(collector, f)
    else:
        print('new plot')
        with open('cache\collector.pkl', 'rb') as f:
            collector = pickle.load(f)
    return collector


if __name__ == '__main__':
    # 切换日期，需要重新生成数据
    # 如果选择了llm，从豆包模型返回，比较慢
    # 如果想要快速输出图片，chat = False
    all_data = load_plot_and_chat_data(
        load_cache=False, 
        save_cache=True, 
        chat=True
        )

    